package com.cloud.ratelimiter;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.cloud.gateway.filter.ratelimit.AbstractRateLimiter;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.validation.Validator;
import org.springframework.validation.annotation.Validated;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.validation.constraints.Min;
import javax.validation.constraints.NotNull;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @Author: TongRui乀
 * @Date: 2020-02-25 11:41
 * @description：  自定义限流器的实现
 */
@Slf4j
public class GatewayRedisRateLimiter extends AbstractRateLimiter<GatewayRedisRateLimiter.Config> implements ApplicationContextAware {

    public static final String CONFIGURATION_PROPERTY_NAME = "custom-redis-rate-limiter";
    public static final String REDIS_SCRIPT_NAME = "redisRequestRateLimiterScript";
    public static final String REMAINING_HEADER = "X-RateLimit-Remaining";
    public static final String REPLENISH_RATE_HEADER = "X-RateLimit-Replenish-Rate";
    public static final String BURST_CAPACITY_HEADER = "X-RateLimit-Burst-Capacity";

    /**
     *  Redis 操作模板
     */
    private ReactiveRedisTemplate<String, String> redisTemplate;
    /**
     *  Redis 脚本
     */
    private RedisScript<List<Long>> script;
    /**
     *  初始化标识
     */
    private AtomicBoolean initialized = new AtomicBoolean(false);
    /**
     *  默认配置
     */
    private Config defaultConfig;

    public GatewayRedisRateLimiter(ReactiveRedisTemplate<String, String> redisTemplate, RedisScript<List<Long>> script,
                                  Validator validator) {
        super(Config.class, CONFIGURATION_PROPERTY_NAME, validator);
        this.redisTemplate = redisTemplate;
        this.script = script;
        initialized.compareAndSet(false, true);
    }

    public GatewayRedisRateLimiter(int defaultReplenishRate, int defaultBurstCapacity) {
        super(Config.class, CONFIGURATION_PROPERTY_NAME, null);
        this.defaultConfig = new Config().setReplenishRate(defaultReplenishRate).setBurstCapacity(defaultBurstCapacity);
    }

    @Override
    public Mono<Response> isAllowed(String routeId, String id) {

        if (!this.initialized.get()) {
            throw new IllegalStateException("RedisRateLimiter is not initialized");
        }

        // 获取限流值  可以通过其他方式动态获取
        GatewayRedisRateLimiter.Config routeConfig = getRouteConfig(routeId);

        int replenishRate = routeConfig.getReplenishRate();
        int burstCapacity = routeConfig.getBurstCapacity();

        try {
            List<String> keys = getKeys(id);
            List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "", Instant.now().getEpochSecond() + "", "1");
            Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);
            return flux.onErrorResume((throwable) -> {
                return Flux.just(Arrays.asList(1L, -1L));
            }).reduce(new ArrayList(), (longs, l) -> {
                longs.addAll(l);
                return longs;
            }).map((results) -> {
                boolean allowed = (Long)results.get(0) == 1L;
                Long tokensLeft = (Long)results.get(1);
                Response response = new Response(allowed, this.getHeaders(replenishRate, burstCapacity, tokensLeft));
                if (this.log.isDebugEnabled()) {
                    this.log.debug("response: " + response);
                }

                return response;
            });
        } catch (Exception var9) {
            this.log.error("Error determining if user allowed from redis", var9);
            return Mono.just(new Response(true, this.getHeaders(replenishRate, burstCapacity, -1L)));
        }
    }

    /**
     *  根据路由标识获取对应的限流值
     *  这里可以实现的功能：  1 动态配置限流  2 同一资源多个限流策略
     * @param routeId
     * @return
     */
    private GatewayRedisRateLimiter.Config getRouteConfig(String routeId) {



        // 从redis 中获取 路由定义 分析出限流值
        Config config = new Config();

        // 若当前路由中没有定义限流 查询全局限流配置
        config.setReplenishRate(15);
        config.setBurstCapacity(15);


        return config;
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {

    }

    @Validated
    public static class Config {
        @Min(1L)
        private int replenishRate;
        @Min(1L)
        private int burstCapacity = 1;

        public Config() {
        }

        public int getReplenishRate() {
            return this.replenishRate;
        }

        public GatewayRedisRateLimiter.Config setReplenishRate(int replenishRate) {
            this.replenishRate = replenishRate;
            return this;
        }

        public int getBurstCapacity() {
            return this.burstCapacity;
        }

        public GatewayRedisRateLimiter.Config setBurstCapacity(int burstCapacity) {
            this.burstCapacity = burstCapacity;
            return this;
        }

        public String toString() {
            return "Config{replenishRate=" + this.replenishRate + ", burstCapacity=" + this.burstCapacity + '}';
        }
    }


    static List<String> getKeys(String id) {
        String prefix = "request_rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    public HashMap<String, String> getHeaders(Integer replenishRate, Integer burstCapacity, Long tokensLeft) {
        HashMap<String, String> headers = new HashMap<>();
        headers.put(this.REMAINING_HEADER, tokensLeft.toString());
        headers.put(this.REPLENISH_RATE_HEADER, String.valueOf(replenishRate));
        headers.put(this.BURST_CAPACITY_HEADER, String.valueOf(burstCapacity));
        return headers;
    }

}
